#include <benchmark/benchmark.h>
#include <thread>
#include <vector>
#include <atomic>
#include <chrono>
#include <string>
#include <utility>
#include <iostream>
#include "blocking_queue.h"
#include "simple_barrier.h"

template<typename T>
struct ValueGenerator {
    static T generate(int producer_id, long long count) {
        return static_cast<T>(producer_id * 1000000 + count);
    }
};

template<>
struct ValueGenerator<std::string> {
    static std::string generate(int producer_id, long long count) {
        return "message_" + std::to_string(producer_id) + "_" +
               std::to_string(count) + "_test_string_data";
    }
};

template<typename T>
static void BM_QueueThroughput(benchmark::State& state)
{
    const int num_producers = state.range(0);
    const int num_consumers = state.range(1);
    const size_t total_op = 100000 * (num_producers + num_consumers);
    const size_t produce_op_per_thread = total_op / num_producers;
    const size_t consume_op_per_thread = total_op / num_consumers;
    const size_t timeout_ms = 5;

    for (auto _ : state) {
        auto queue = BlockingQueue<T>::create(10000);
        std::atomic<bool> stop_flag{false};

        std::vector<std::thread> producers;
        std::vector<std::thread> consumers;
        std::atomic<long long> total_operations{0};
        std::atomic<long long> success_operations{0};
        std::atomic<long long> failed_operations{0};

        SimpleBarrier start_barrier(num_producers + num_consumers + 1);

        for (int i = 0; i < num_producers; ++i) {
            producers.emplace_back([&, i]() {
                start_barrier.wait();

                long long local_count = 0;
                long long local_success = 0;
                long long local_failed = 0;

                while (local_count < produce_op_per_thread && !stop_flag.load()) {
                    T value = ValueGenerator<T>::generate(i, local_count);

                    if (queue->push(std::move(value), timeout_ms)) {
                        local_success++;
                    } else {
                        local_failed++;
                    }
                    local_count++;
                }

                total_operations.fetch_add(local_count, std::memory_order_relaxed);
                success_operations.fetch_add(local_success, std::memory_order_relaxed);
                failed_operations.fetch_add(local_failed, std::memory_order_relaxed);
            });
        }

        for (int i = 0; i < num_consumers; ++i) {
            consumers.emplace_back([&]() {
                start_barrier.wait();

                T value;
                long long local_count = 0;
                long long local_success = 0;
                long long local_failed = 0;

                while (local_count < consume_op_per_thread && !stop_flag.load()) {
                    if (queue->pop(value, timeout_ms)) {
                        local_success++;
                    } else {
                        local_failed++;
                    }
                    local_count++;
                }
                total_operations.fetch_add(local_count, std::memory_order_relaxed);
                success_operations.fetch_add(local_success, std::memory_order_relaxed);
                failed_operations.fetch_add(local_failed, std::memory_order_relaxed);
            });
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        auto start_time = std::chrono::steady_clock::now();
        start_barrier.wait();

        for (auto& producer : producers) {
            if (producer.joinable()) producer.join();
        }
        for (auto& consumer : consumers) {
            if (consumer.joinable()) consumer.join();
        }

        auto end_time = std::chrono::steady_clock::now();
        double duration_seconds = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time).count();
        long long total_success = success_operations.load();
        double throughput = total_success / duration_seconds;

        state.counters["Through"] = throughput;
        state.counters["TotalOps"] = total_operations.load();
        state.counters["FailedOps"] = failed_operations.load();
        state.counters["Duration"] = duration_seconds;

        state.SetItemsProcessed(total_success);
    }
}

struct CustomData {
    int id;
    double value;
    std::string name;

    CustomData(int i = 0, double v = 0.0, std::string n = "")
        : id(i), value(v), name(std::move(n)) {}

    CustomData(CustomData&& other) noexcept
        : id(other.id), value(other.value), name(std::move(other.name)) {}

    CustomData& operator=(CustomData&& other) noexcept {
        if (this != &other) {
            id = other.id;
            value = other.value;
            name = std::move(other.name);
        }
        return *this;
    }

    CustomData(const CustomData&) = delete;
    CustomData& operator=(const CustomData&) = delete;
};

template<>
struct ValueGenerator<CustomData> {
    static CustomData generate(int producer_id, long long count) {
        return CustomData(
            producer_id,
            count * 1.5,
            "custom_" + std::to_string(producer_id) + "_" + std::to_string(count)
        );
    }
};

#define REGISTER_BENCHMARK_FOR_TYPE(FuncName, TypeName) \
    static void FuncName(benchmark::State& state) { \
        BM_QueueThroughput<TypeName>(state); \
    } \
    BENCHMARK(FuncName)->Args({1, 1})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({2, 2})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({4, 4})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({8, 8})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({16, 16})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({2, 1})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({4, 1})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({8, 1})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({16, 1})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({1, 2})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({1, 4})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({1, 8})->UseRealTime(); \
    BENCHMARK(FuncName)->Args({1, 16})->UseRealTime();


REGISTER_BENCHMARK_FOR_TYPE(BM_IntAccurateThroughput, int)
REGISTER_BENCHMARK_FOR_TYPE(BM_StrAccurateThroughput, std::string)
REGISTER_BENCHMARK_FOR_TYPE(BM_LongAccurateThroughput, long long)
REGISTER_BENCHMARK_FOR_TYPE(BM_CustomAccurateThroughput, CustomData)
